[Feat] Enable VAE parallel in HunyuanImage3#3091
Merged
Gaohan123 merged 8 commits intoMay 20, 2026
Merged
Conversation
421d557 to
c69899e
Compare
Contributor
Author
ee9b0b3 to
a4502c4
Compare
Collaborator
|
does it work in GPU as well? |
a4fc4ec to
378289a
Compare
1 task
378289a to
2eacaf2
Compare
ae99ecc to
c6f0e06
Compare
Contributor
c6f0e06 to
8c4b866
Compare
Gaohan123
reviewed
May 14, 2026
Collaborator
Gaohan123
left a comment
There was a problem hiding this comment.
Here are some suggestions:
- Please add simple UT for it
- I didn't notice any modification about NPU, which is not consistent with your PR description
Contributor
Author
|
f338f4d to
014b54b
Compare
Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Signed-off-by: KexiongYu <yukexiong1@huawei.com>
014b54b to
272bc98
Compare
…oencoder_kl_hunyuan.py to clarify what each case validates, without changing test behavior Signed-off-by: zzh <943967662@qq.com>
4f1aa1b to
61fc92a
Compare
RuixiangMa
reviewed
May 19, 2026
| logger = init_logger(__name__) | ||
|
|
||
|
|
||
| class DistributedAutoencoderKLHunyuan(AutoencoderKLConv3D, DistributedVaeMixin): |
Collaborator
There was a problem hiding this comment.
missing from_pretrained, consistency suggests adding it
| def encode_tile_exec(self, task: TileTask) -> torch.Tensor: | ||
| return self.encoder(task.tensor) | ||
|
|
||
| def encode_tile_merge( |
Collaborator
There was a problem hiding this comment.
tile_merge and encode_tile_merge are byte-for-byte identical. Could extract helper
| torch.nn.Module.__init__(self) | ||
| self.tile_latent_min_size = 2 | ||
| self.tile_sample_min_size = 2 | ||
| self.tile_overlap_factor = 0.0 |
Collaborator
There was a problem hiding this comment.
tile_overlap_factor hardcoded to 0.0 (real default 0.25), blend logic never tested.
autoencoder_kl_hunyuan imports AutoencoderKLConv3D from the hunyuan_image3 package, which triggers hunyuan_image3/__init__.py to execute and import pipeline_hunyuan_image3, which in turn imported DistributedAutoencoderKLHunyuan back from autoencoder_kl_hunyuan before it finished initializing, causing a circular import error during test collection. Fix by moving the top-level import of DistributedAutoencoderKLHunyuan into HunyuanImage3Pipeline.__init__ as a lazy import, so it is only resolved at call time when both modules are fully initialized. Signed-off-by: zzh <943967662@qq.com>
autoencoder_kl_hunyuan imports AutoencoderKLConv3D from the hunyuan_image3 package, which triggers hunyuan_image3/__init__.py to execute and import pipeline_hunyuan_image3, which in turn imported DistributedAutoencoderKLHunyuan back from autoencoder_kl_hunyuan before it finished initializing, causing a circular import error during test collection. Fix by moving the top-level import of DistributedAutoencoderKLHunyuan into HunyuanImage3Pipeline.__init__ as a lazy import, so it is only resolved at call time when both modules are fully initialized. Signed-off-by: zzh <943967662@qq.com>
…an and deduplicate tile merge - Add missing from_pretrained classmethod for consistency with other distributed autoencoders (KL, Wan, QwenImage) - Delegate encode_tile_merge to tile_merge to eliminate byte-for-byte duplicate code Signed-off-by: zzh <943967662@qq.com>
…AE tests Adjust grid_shape from (2,2) to (4,4) and tile count from 4 to 16. When tile_overlap_factor=0.25, overlap_size becomes 1 instead of 2, producing a denser 4x4 tile grid on the 4x4 input. Signed-off-by: zzh <943967662@qq.com>
…AE tests With tile_latent_min_size=2 and tile_overlap_factor=0.25, blend_extent truncates to int(0.5)=0, causing overlapping tiles with no blending and producing misaligned 7x7 output instead of the expected 4x4. Increasing min_size to 8 makes blend_extent=2 and keeps the tile pipeline's math self-consistent while preserving tile_overlap_factor=0.25. Signed-off-by: zzh <943967662@qq.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.


Summary
Enable VAE parallel support in HunyuanImage3.
Current changes:
vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.pyHunyuanImage3Pipelineto use the distributed autoencoder wrappervllm_omni/platforms/npu/models/hunyuan_fused_moe.pyunified deploy yaml in #3172
Validation
py_compile, diff checks)Test Plan
Tested on 4xAscend NPU
server
vae_patch_parallel_size is set to 4
client
Test Result
output
VAE decode time 625.7ms -> 355ms
w/o vae parallel

w vae parallel
